"""
MCP客户端 - 基于LangGraph的MCP适配器
"""

import asyncio
from typing import Dict, List, Any, Optional

from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_mcp_adapters.sessions import Connection
from langchain_mcp_adapters.tools import load_mcp_tools
from pydantic import BaseModel

from config.logging_config import get_logger

logger = get_logger(__name__)


class MCPServerConfig(BaseModel):
    """MCP服务器配置"""
    name: str
    transport: str  # "stdio", "streamable_http", "sse", "websocket"
    command: Optional[str] = None
    args: Optional[List[str]] = None
    url: Optional[str] = None
    env: Optional[Dict[str, str]] = None


class MCPToolInfo(BaseModel):
    """MCP工具信息"""
    name: str
    description: str
    input_schema: Dict[str, Any]
    server_name: str


class MCPClient:
    """MCP客户端管理类"""

    def __init__(self):
        self.client: Optional[MultiServerMCPClient] = None
        self.server_configs: Dict[str, MCPServerConfig] = {}
        self.connections: Dict[str, Connection] = {}
        self.available_tools: Dict[str, MCPToolInfo] = {}

    def _create_connection(self, config: MCPServerConfig) -> Connection:
        """根据配置创建连接对象"""
        if config.transport == "stdio":
            if not config.command:
                raise ValueError("stdio transport requires command")

            connection: Connection = {
                "transport": "stdio",
                "command": config.command,
                "args": config.args or [],
            }
            if config.env:
                connection["env"] = config.env

        elif config.transport == "streamable_http":
            if not config.url:
                raise ValueError("streamable_http transport requires url")

            connection = {
                "transport": "streamable_http",
                "url": config.url,
            }

        elif config.transport == "sse":
            if not config.url:
                raise ValueError("sse transport requires url")

            connection = {
                "transport": "sse",
                "url": config.url,
            }

            if config.env:
                connection["env"] = config.env

            logger.info(f"Creating SSE connection to: {config.url}")

        elif config.transport == "websocket":
            if not config.url:
                raise ValueError("websocket transport requires url")

            connection = {
                "transport": "websocket",
                "url": config.url,
            }
        else:
            raise ValueError(f"Unsupported transport type: {config.transport}")

        return connection

    def _rebuild_client(self):
        """重新构建MultiServerMCPClient"""
        if self.connections:
            self.client = MultiServerMCPClient(self.connections)
        else:
            self.client = None

    async def _execute_with_session(self, server_name: str, operation):
        """在会话中执行操作，处理可能的异常"""
        try:
            async with self.client.session(server_name) as session:
                return await operation(session)
        except Exception as e:
            # 检查是否是连接相关的异常，进行一次重试
            error_msg = str(e)
            if any(phrase in error_msg.lower() for phrase in [
                "connection", "timeout", "network", "http"
            ]):
                logger.warning(f"连接异常，将重试: {server_name} - {error_msg}")
                await asyncio.sleep(0.5)
                try:
                    async with self.client.session(server_name) as session:
                        return await operation(session)
                except Exception as retry_error:
                    logger.error(f"重试失败: {server_name} - {str(retry_error)}")
                    raise retry_error
            else:
                raise e

    async def connect_server(self, config: MCPServerConfig) -> bool:
        """连接单个MCP服务器"""
        try:
            logger.info(f"连接MCP服务器: {config.name} ({config.transport})")

            # 创建连接配置
            connection = self._create_connection(config)

            # 添加到连接字典
            self.connections[config.name] = connection
            self.server_configs[config.name] = config

            # 重新构建客户端
            self._rebuild_client()

            # 测试连接
            if self.client:
                try:
                    async def test_operation(session):
                        tools = await load_mcp_tools(session)
                        return tools

                    tools = await self._execute_with_session(config.name, test_operation)
                    logger.info(f"MCP服务器连接成功: {config.name}, 工具数量: {len(tools)}")

                except Exception as conn_error:
                    logger.error(f"MCP服务器连接失败: {config.name} - {str(conn_error)}")

                    # 连接失败，回滚更改
                    self.connections.pop(config.name, None)
                    self.server_configs.pop(config.name, None)
                    self._rebuild_client()
                    return False

            return True

        except Exception as e:
            logger.error(f"连接MCP服务器异常: {config.name} - {str(e)}")
            # 清理失败的连接
            self.connections.pop(config.name, None)
            self.server_configs.pop(config.name, None)
            self._rebuild_client()
            return False

    async def connect_from_json_schema(self, server_name: str, json_schema: Dict[str, Any]) -> bool:
        """从JSON配置连接MCP服务器"""
        try:
            # 解析JSON配置为MCPServerConfig
            transport = json_schema.get("transport")
            if not transport:
                if "command" in json_schema:
                    transport = "stdio"
                elif json_schema.get("url") or json_schema.get("endpoint"):
                    transport = "streamable_http"
                else:
                    transport = "streamable_http"

            # 获取URL，支持url和endpoint两种字段名
            url = json_schema.get("url") or json_schema.get("endpoint")

            # 创建配置对象
            config = MCPServerConfig(
                name=server_name,
                transport=transport,
                command=json_schema.get("command"),
                args=json_schema.get("args", []),
                url=url,
                env=json_schema.get("env")
            )

            # 验证配置完整性
            if config.transport == "stdio" and not config.command:
                raise ValueError(f"stdio传输需要command字段")
            elif config.transport in ["streamable_http", "sse", "websocket"] and not config.url:
                raise ValueError(f"{config.transport}传输需要url或endpoint字段")

            logger.info(f"从JSON配置连接MCP服务器: {server_name}")

            # 尝试连接
            result = await self.connect_server(config)
            if not result:
                logger.error(f"JSON配置连接失败: {server_name}")
            return result

        except Exception as e:
            logger.error(f"解析JSON配置失败: {server_name} - {str(e)}")
            return False

    async def get_tools(self, server_name: Optional[str] = None) -> List[MCPToolInfo]:
        """获取工具列表"""
        try:
            if self.client is None:
                logger.warning("MCP客户端未连接")
                return []

            tool_list = []

            if server_name:
                # 获取特定服务器的工具
                if server_name not in self.connections:
                    logger.warning(f"服务器未连接: {server_name}")
                    return []

                try:
                    async def get_tools_operation(session):
                        tools = await load_mcp_tools(session)
                        return tools

                    tools = await self._execute_with_session(server_name, get_tools_operation)

                    for tool in tools:
                        input_schema = {}
                        try:
                            if hasattr(tool, 'args_schema') and tool.args_schema:
                                if hasattr(tool.args_schema, 'model_json_schema'):
                                    input_schema = tool.args_schema.model_json_schema()
                                elif hasattr(tool.args_schema, 'schema'):
                                    input_schema = tool.args_schema.schema()
                                elif isinstance(tool.args_schema, dict):
                                    input_schema = tool.args_schema
                        except Exception:
                            input_schema = {}

                        tool_info = MCPToolInfo(
                            name=tool.name,
                            description=tool.description or "",
                            input_schema=input_schema,
                            server_name=server_name
                        )
                        tool_list.append(tool_info)
                        self.available_tools[tool.name] = tool_info

                except Exception as e:
                    logger.error(f"获取服务器工具失败: {server_name} - {str(e)}")
            else:
                # 获取所有服务器的工具
                for connected_server in self.connections.keys():
                    try:
                        async def get_tools_operation(session):
                            tools = await load_mcp_tools(session)
                            return tools

                        tools = await self._execute_with_session(connected_server, get_tools_operation)

                        for tool in tools:
                            input_schema = {}
                            try:
                                if hasattr(tool, 'args_schema') and tool.args_schema:
                                    if hasattr(tool.args_schema, 'model_json_schema'):
                                        input_schema = tool.args_schema.model_json_schema()
                                    elif hasattr(tool.args_schema, 'schema'):
                                        input_schema = tool.args_schema.schema()
                                    elif isinstance(tool.args_schema, dict):
                                        input_schema = tool.args_schema
                            except Exception:
                                input_schema = {}

                            tool_info = MCPToolInfo(
                                name=tool.name,
                                description=tool.description or "",
                                input_schema=input_schema,
                                server_name=connected_server
                            )
                            tool_list.append(tool_info)
                            self.available_tools[tool.name] = tool_info

                    except Exception as e:
                        logger.error(f"获取服务器工具失败: {connected_server} - {str(e)}")

            logger.info(f"获取到 {len(tool_list)} 个MCP工具")
            return tool_list

        except Exception as e:
            logger.error(f"获取MCP工具失败: {str(e)}")
            return []

    async def disconnect_server(self, server_name: str) -> bool:
        """断开MCP服务器连接"""
        try:
            if server_name not in self.connections:
                logger.warning(f"服务器未连接: {server_name}")
                return False

            logger.info(f"断开MCP服务器: {server_name}")

            # 移除连接和配置
            self.connections.pop(server_name, None)
            self.server_configs.pop(server_name, None)

            # 清理相关工具
            tools_to_remove = [name for name, tool in self.available_tools.items()
                               if tool.server_name == server_name]
            for tool_name in tools_to_remove:
                self.available_tools.pop(tool_name, None)

            # 重新构建客户端
            self._rebuild_client()

            logger.info(f"MCP服务器断开成功: {server_name}")
            return True

        except Exception as e:
            logger.error(f"断开MCP服务器失败: {server_name} - {str(e)}")

            # 强制清理状态
            self.connections.pop(server_name, None)
            self.server_configs.pop(server_name, None)
            tools_to_remove = [name for name, tool in self.available_tools.items()
                               if tool.server_name == server_name]
            for tool_name in tools_to_remove:
                self.available_tools.pop(tool_name, None)
            self._rebuild_client()

            return False

    async def test_connection(self, server_name: str) -> bool:
        """测试服务器连接"""
        try:
            if self.client is None or server_name not in self.connections:
                return False

            async def test_operation(session):
                tools = await load_mcp_tools(session)
                return tools

            tools = await self._execute_with_session(server_name, test_operation)
            logger.debug(f"连接测试成功: {server_name}, 工具数量: {len(tools)}")
            return True

        except Exception as e:
            logger.warning(f"连接测试失败: {server_name} - {str(e)}")
            return False

    def get_connected_servers(self) -> List[str]:
        """获取已连接的服务器列表"""
        return list(self.connections.keys())

    def get_tool_by_name(self, tool_name: str) -> Optional[MCPToolInfo]:
        """根据名称获取工具信息"""
        return self.available_tools.get(tool_name)

    async def get_langchain_tools(self, server_name: Optional[str] = None) -> List[Any]:
        """获取LangChain工具对象（用于直接在LangGraph中使用）"""
        try:
            if self.client is None:
                logger.warning("MCP客户端未连接")
                return []

            if server_name:
                return await self.client.get_tools(server_name=server_name)
            else:
                return await self.client.get_tools()

        except Exception as e:
            logger.error(f"获取LangChain工具失败: {str(e)}")
            return []

    async def call_tool(self, server_name: Optional[str], tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """使用LangGraph标准方式调用MCP工具"""
        try:
            if self.client is None:
                return {"error": "MCP客户端未连接", "type": "client_error"}

            logger.info(f"调用MCP工具: {tool_name}")

            # 如果指定了服务器名称，直接使用
            if server_name and server_name in self.connections:
                return await self.call_tool_by_server(server_name, tool_name, arguments)

            # 自动检测工具所在的服务器
            target_server = None
            target_tool = None

            for connected_server in self.connections.keys():
                try:
                    async def find_tool_operation(session):
                        tools = await load_mcp_tools(session)
                        for tool in tools:
                            if tool.name == tool_name:
                                return tool
                        return None

                    found_tool = await self._execute_with_session(connected_server, find_tool_operation)
                    if found_tool:
                        target_server = connected_server
                        target_tool = found_tool
                        break

                except Exception:
                    continue

            if not target_tool or not target_server:
                available_tools = []
                for server in self.connections.keys():
                    try:
                        async def get_tool_names_operation(session):
                            tools = await load_mcp_tools(session)
                            return [t.name for t in tools]

                        tool_names = await self._execute_with_session(server, get_tool_names_operation)
                        available_tools.extend(tool_names)
                    except:
                        pass

                return {
                    "error": f"未找到工具: {tool_name}",
                    "type": "tool_not_found",
                    "available_tools": available_tools
                }

            # 调用工具
            try:
                result = await target_tool.ainvoke(arguments)
                logger.info(f"MCP工具调用成功: {tool_name}")
                return {
                    "success": True,
                    "tool_name": tool_name,
                    "arguments": arguments,
                    "result": result,
                    "server_name": target_server
                }

            except Exception as tool_error:
                logger.error(f"MCP工具执行失败: {tool_name} - {str(tool_error)}")
                return {
                    "error": f"工具执行失败: {str(tool_error)}",
                    "type": "tool_execution_error",
                    "tool_name": tool_name
                }

        except Exception as e:
            logger.error(f"MCP工具调用异常: {tool_name} - {str(e)}")
            return {
                "error": f"工具调用异常: {str(e)}",
                "type": "call_error",
                "tool_name": tool_name
            }

    def call_tool_sync(self, server_name: Optional[str], tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """同步版本的工具调用（为了兼容性）"""
        try:
            # 如果当前已在事件循环中，创建新的事件循环
            try:
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    # 在已运行的事件循环中，使用run_in_executor
                    import concurrent.futures
                    with concurrent.futures.ThreadPoolExecutor() as executor:
                        future = executor.submit(
                            lambda: asyncio.run(self.call_tool(server_name, tool_name, arguments))
                        )
                        return future.result(timeout=30)
                else:
                    return loop.run_until_complete(self.call_tool(server_name, tool_name, arguments))
            except RuntimeError:
                # 没有事件循环，创建新的
                return asyncio.run(self.call_tool(server_name, tool_name, arguments))

        except Exception as e:
            logger.error(f"同步MCP工具调用失败: {tool_name} - {str(e)}")
            return {
                "error": f"同步调用失败: {str(e)}",
                "type": "sync_call_error",
                "tool_name": tool_name
            }

    async def call_tool_by_server(self, server_name: str, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """通过指定服务器调用工具"""
        try:
            if self.client is None:
                return {"error": "MCP客户端未连接", "type": "client_error"}

            if server_name not in self.connections:
                return {
                    "error": f"服务器未连接: {server_name}",
                    "type": "server_not_connected"
                }

            async def call_tool_operation(session):
                tools = await load_mcp_tools(session)

                # 查找指定工具
                target_tool = None
                for tool in tools:
                    if tool.name == tool_name:
                        target_tool = tool
                        break

                if not target_tool:
                    return {
                        "error": f"服务器中未找到工具: {tool_name}",
                        "type": "tool_not_found",
                        "available_tools": [t.name for t in tools]
                    }

                # 调用工具
                result = await target_tool.ainvoke(arguments)
                return {
                    "success": True,
                    "tool_name": tool_name,
                    "arguments": arguments,
                    "result": result,
                    "server_name": server_name
                }

            result = await self._execute_with_session(server_name, call_tool_operation)

            if result.get("success"):
                logger.info(f"服务器工具调用成功: {server_name}.{tool_name}")

            return result

        except Exception as e:
            logger.error(f"服务器工具调用失败: {server_name}.{tool_name} - {str(e)}")
            return {
                "error": f"服务器调用失败: {str(e)}",
                "type": "server_call_error",
                "tool_name": tool_name
            }

    async def get_tool_schema(self, tool_name: str, server_name: Optional[str] = None) -> Dict[str, Any]:
        """获取工具的参数模式"""
        try:
            if self.client is None:
                return {"error": "MCP客户端未连接"}

            # 获取工具
            if server_name:
                tools = await self.client.get_tools(server_name=server_name)
            else:
                tools = await self.client.get_tools()

            for tool in tools:
                if tool.name == tool_name:
                    schema = {}
                    try:
                        if hasattr(tool, 'args_schema') and tool.args_schema:
                            if hasattr(tool.args_schema, 'model_json_schema'):
                                schema = tool.args_schema.model_json_schema()
                            elif hasattr(tool.args_schema, 'schema'):
                                schema = tool.args_schema.schema()
                    except Exception:
                        pass

                    return {
                        "tool_name": tool_name,
                        "description": tool.description or "",
                        "schema": schema,
                        "server_name": getattr(tool, '_server_name', server_name or 'unknown')
                    }

            return {"error": f"未找到工具: {tool_name}"}

        except Exception as e:
            logger.error(f"获取工具模式失败: {tool_name} - {str(e)}")
            return {"error": f"获取工具模式失败: {str(e)}"}

    async def cleanup(self):
        """清理资源"""
        try:
            logger.info("开始清理MCPClient资源...")

            # 断开所有已连接的服务器
            if self.client is not None:
                connected_servers = list(self.connections.keys())
                if connected_servers:
                    logger.info(f"断开 {len(connected_servers)} 个MCP服务器连接")

                for server_name in connected_servers:
                    try:
                        await self.disconnect_server(server_name)
                    except Exception:
                        pass

                # 尝试正常关闭客户端
                try:
                    if hasattr(self.client, '__aexit__'):
                        await self.client.__aexit__(None, None, None)
                except Exception as e:
                    logger.debug(f"客户端关闭异常: {str(e)}")

            # 清理内部状态
            self.connections.clear()
            self.server_configs.clear()
            self.available_tools.clear()
            self.client = None

            # 等待异步任务完成
            await asyncio.sleep(0.3)

            logger.info("MCPClient资源清理完成")

        except Exception as e:
            logger.error(f"MCPClient清理异常: {str(e)}")

            # 强制清理内部状态
            try:
                self.connections.clear()
                self.server_configs.clear()
                self.available_tools.clear()
                self.client = None
            except Exception:
                pass


# 全局MCP客户端实例
mcp_client = MCPClient()
